package com.xbongbong.paas.gateway.decorator;

import io.netty.buffer.UnpooledByteBufAllocator;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import reactor.core.publisher.Flux;

import java.nio.ByteBuffer;

import static reactor.core.scheduler.Schedulers.single;

/**
 * 用以解决request.getBody只能读一次的问题
 * @author kaka
 * @time 2018-8-3 19:09
 */
public class CacheBodyServerHttpRequestDecorator extends ServerHttpRequestDecorator {

    private int getBufferTime = 0;
    private byte[] bytes;

    public CacheBodyServerHttpRequestDecorator(ServerHttpRequest delegate) {
        super(delegate);
    }

    @Override
    public Flux<DataBuffer> getBody() {
        if (getBufferTime == 0) {
            getBufferTime++;
            Flux<DataBuffer> flux = super.getBody();
            return flux
                    .publishOn(single())
                    .map(this::cache);
        } else {
            return Flux.just(getBodyMore());
        }
    }

    private DataBuffer getBodyMore() {
        NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(new UnpooledByteBufAllocator(false));
        //空body处理
        if(bytes == null) {
            bytes = new byte[0];
        }

        return nettyDataBufferFactory.wrap(bytes);
    }

    private DataBuffer cache(DataBuffer buffer) {

        ByteBuffer tmpByteBuffer = buffer.asByteBuffer();
        byte[] tmpBytes = tmpByteBuffer.array();

        //合并byte
        if(bytes != null) {
            bytes = unitByteArray(bytes, tmpBytes);
        } else {
            bytes = tmpBytes;
        }

        NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(new UnpooledByteBufAllocator(false));
        return nettyDataBufferFactory.wrap(tmpBytes);
    }

    /**
     * 合并byte数组
     */
    private static byte[] unitByteArray(byte[] byte1, byte[] byte2){
        byte[] unitByte = new byte[byte1.length + byte2.length];
        System.arraycopy(byte1, 0, unitByte, 0, byte1.length);
        System.arraycopy(byte2, 0, unitByte, byte1.length, byte2.length);
        return unitByte;
    }
}
